import datetime
import os
import sys
from concurrent.futures import ProcessPoolExecutor, as_completed
from decimal import Decimal
from typing import *
from executor.base import BaseExecutor
from utils.data.loader import BaseLoader
from utils.data.sqlLoader import loader_mapper
from utils.dataArrange.finalReport import FinalReport
from utils.loggers import Logger
from loguru import logger
from config import ORIGIN_BALANCE, BACKTEST_DATA_COLLECTION, ROOT_DIR, DATABASE_TYPE, STRATEGY_LOG
from collections import defaultdict
import pickle
from utils.dataArrange.dataHandler import handler_map, Handler
import platform
from globalVar import GlobalVar
from utils.error import PreparationError
import subprocess

import pandas as pd

from utils.strategy.configType import ConfigType
from utils.strategy.loader import FactorLoader
from utils.tools import NullStdout
from utils.trade.tradeType import FakeTradeType


class StepExecutor(BaseExecutor):
    """
    单个策略的执行器
    """

    def __init__(self, strategy_name: str):
        super().__init__(strategy_name)
        # 策略自定义保存的二进制数据，初始值为空字典
        self.strategy_cache: bytes = b'\x80\x04}\x94.'
        self.strategies_result: pd.DataFrame = pd.DataFrame()
        # 交易对象
        self.traders: Optional[FakeTradeType] = None
        # 临时数据
        self.template: Dict[str, Any] = {}

    def initial_attrs(self, conf: ConfigType = None):
        """
        顺序加载成员变量
        :param conf:
        :return:
        """
        super().initial_attrs(conf)
        # 加载交易对象
        self.traders = FakeTradeType(self.strategy.arg_id)
        self.loader: BaseLoader = loader_mapper.get(DATABASE_TYPE, BaseLoader)(
            start=self.s_time - self.max_line_count * self.period * self.data_batch, end=self.e_time,
            period=self.period, symbols=(self.symbols + ['BTCUSDT']) if self.symbols else None,
            with_decimal=True, batch_size=self.data_batch)
        self.dataframe = self.loader.fetch(self.s_time - self.period * self.max_line_count * self.data_batch,
                                           self.max_line_count - 1)
        # 数据准备完成后，移动cursor
        self.loader.cursor = self.s_time - self.period * self.data_batch
        GlobalVar.EXE_TIMESTAMP = self.s_time

        # stdout文件
        self.logfile = open(
            os.path.join(ROOT_DIR, f'locals/output', f'stdout-{self.strategy.name}.log'), 'w+',
            encoding='utf-8') if STRATEGY_LOG else NullStdout() if not self.async_type and STRATEGY_LOG else NullStdout

    @staticmethod
    def update_price(data: pd.DataFrame):
        """

        :param data:
        :return:
        """
        # 更新币种最新价格为当前K线数据
        spot = data[data['trade_type'] == 0]
        GlobalVar.SPOT_PRICE |= {x: y for x, y in zip(spot['symbol'], spot['close'])}
        swap = data[data['trade_type'] == 1]
        GlobalVar.SWAP_PRICE |= {x: y for x, y in zip(swap['symbol'], swap['close'])}

    def step_data(self):
        """

        :return:
        """
        # 重置策略结果
        self.strategies_result = pd.DataFrame()
        # 步进K线数据
        new_data: pd.DataFrame = self.loader.read()
        self.update_price(new_data)
        self.dataframe = pd.concat([self.dataframe, new_data], ignore_index=True)
        self.dataframe = self.dataframe[self.dataframe[self.loader.date_column] >= (
                GlobalVar.EXE_TIMESTAMP - self.max_line_count * self.period) * 1000]

    def step(self):
        """

        :return:
        """
        self.step_data()
        sys.stdout = self.logfile
        # 运行交易对象步进函数
        self.traders.exchange.step()
        # 注入运行时间到策略中
        self.strategy.injection({'execute_timestamp': GlobalVar.EXE_TIMESTAMP})
        if self.strategy.is_execute():
            # 注入因子逻辑
            self.strategy.injection({'factors': {k: FactorLoader(k).signal() for k in self.strategy.factors}})
            self.strategies_result = self.strategy.pick(
                self.dataframe[
                    self.dataframe['symbol'].isin(self.symbols)].copy() if self.symbols else self.dataframe.copy())
            if self.is_backtest and BACKTEST_DATA_COLLECTION:
                bt_data = self.strategy.backtest_data()
                for k, v in bt_data.items():
                    handler: Handler = handler_map[k](self.strategy.name)
                    handler.read(v)
                    if k not in self.backtest_handlers:
                        self.backtest_handlers[k] = handler
            # 注入自定义数据和接口
            self.strategy.injection({'custom': pickle.loads(self.strategy_cache),
                                     'spot_trade': self.traders.spot_trade,
                                     'swap_trade': self.traders.swap_trade,
                                     'transfer': self.traders.transfer,
                                     'margin_type': self.traders.margin_type,
                                     'position_side': self.traders.position_side,
                                     'swap_leverage': self.traders.swap_leverage,
                                     'spot_balance_interface': self.traders.exchange.spot_balance,
                                     'swap_balance_interface': self.traders.exchange.swap_balance,
                                     'spot_price': GlobalVar.SPOT_PRICE.copy(),
                                     'swap_price': GlobalVar.SWAP_PRICE.copy(),
                                     'spot_balance': self.traders.exchange.spot_balance(),
                                     'swap_balance': self.traders.exchange.swap_balance(),
                                     'swap_position': self.traders.exchange.swap_position(),
                                     'is_backtest': self.is_backtest
                                     })
            self.strategy.trade(self.strategies_result)
            # 保存策略的自定义数据
            self.strategy_cache = self.strategy.custom()
        sys.stdout = self.stdout_bak
        self.dump_info()

    def dump_info(self):
        """
        每日0点收集信息
        :return:
        """
        if GlobalVar.EXE_TIMESTAMP % (24 * 3600) != 0:
            return

        last_funding: float = self.template.get('lastFunding', int(ORIGIN_BALANCE))
        funding = round(float(self.traders.exchange.account_balance()), 2)
        temp = {
            'BTC': round(float(GlobalVar.SPOT_PRICE['BTCUSDT']), 2),  # btc价格
            'funding': funding,  # 资产价值
            'tradeSpot': round(float(self.traders.exchange.attr_asset_quantity['spot']), 2),  # 现货交易额
            'tradeSwap': round(float(self.traders.exchange.attr_asset_quantity['swap']), 2),  # 合约交易额
            # 未使用资金
            'unusedFunding': round(float(self.traders.exchange.attr_spot_balance['USDT'] +
                                         self.traders.exchange.attr_swap_balance['USDT']['maxWithdrawAmount']), 2),
            'spotPremium': round(float(self.traders.exchange.attr_premium['spot']), 6),  # 现货手续费
            'swapPremium': round(float(self.traders.exchange.attr_premium['swap']), 6),  # 合约手续费
            'netIncome': round(funding - last_funding, 2),  # 净收益

        }
        self.template['lastFunding'] = funding
        # 重置交易统计信息
        self.traders.exchange.attr_asset_quantity = defaultdict(Decimal)
        self.traders.exchange.attr_premium = defaultdict(Decimal)

        self.result_data[GlobalVar.EXE_TIMESTAMP] = temp

    def draw(self) -> Any:
        """
        绘图、计算指标
        :return:
        """
        fin = FinalReport(f'report-{self.strategy.arg_id}')
        times = [datetime.datetime.fromtimestamp(x, tz=datetime.timezone.utc).strftime('%Y-%m-%d') for x in
                 self.result_data.keys()]
        funding: List[Any] = [x['funding'] for x in self.result_data.values()]

        fin.account_funding(
            times,
            self.strategy.arg_id, funding
            , target=[x['BTC'] for x in self.result_data.values()],
            net_income=[x['netIncome'] for x in self.result_data.values()])
        fin.trade_info(times,
                       [x['tradeSpot'] for x in self.result_data.values()],
                       [x['tradeSwap'] for x in self.result_data.values()],
                       funding,
                       [x['unusedFunding'] for x in self.result_data.values()],
                       [x['spotPremium'] for x in self.result_data.values()],
                       [x['swapPremium'] for x in self.result_data.values()],
                       )

        fin.infos['累计净值'] = round(funding[-1] / funding[0], 2)
        fin.infos['年化'] = round(
            abs(fin.infos['累计净值']) ** (365 / (GlobalVar.END_TIME - GlobalVar.START_TIME).days) - 1, 2)
        fin.infos['月化'] = round((abs(fin.infos['年化']) + 1) ** (1 / 12) - 1, 2) * (
            1 if fin.infos['年化'] >= 0 else -1)
        withdrawals = fin.withdrawal(funding)
        fin.info_table(
            {'最大回撤时间': f'{times[withdrawals[0]]} - {times[withdrawals[1]]}', '最大回撤': withdrawals[2]})

        if not self.async_type:
            fin.generate(os.path.join(ROOT_DIR, 'locals'))
        return fin.infos

    @staticmethod
    def _stepper(ste: 'StepExecutor', conf: ConfigType = None) -> Tuple[str, Dict[str, Any]]:
        """
        单次配置执行
        :param ste:
        :param conf:
        :return:
        """
        ste.initial_attrs(conf)
        logger.info(f'开始计算: {ste.strategy.name}， 参数：{conf.__str__()}')
        while GlobalVar.EXE_TIMESTAMP < ste.e_time:
            ste.step()
            # 单任务时显示时间进度
            if not ste.async_type:
                Logger.info(f'{datetime.datetime.fromtimestamp(GlobalVar.EXE_TIMESTAMP, tz=datetime.timezone.utc)}完成',
                            new_line=False)
            # 时间步进
            GlobalVar.EXE_TIMESTAMP += ste.period * ste.data_batch
        print('\n' if not ste.async_type else '')
        logger.info(f'计算完成: {ste.strategy.name}， 参数：{conf.__str__()}')
        logger.info(f'数据整理{ste.strategy.arg_id}')
        infos = ste.draw()
        if ste.is_backtest:
            for k, v in ste.backtest_handlers.items():
                v.dump(os.path.join(ROOT_DIR, 'locals'))
        logger.info(f'数据整理完成')

        html_f_name = f'report-{ste.strategy.arg_id}' + '.html'
        if platform.system() == 'Windows' and not ste.async_type:
            os.system(fr'start {os.path.join(ROOT_DIR, "locals", html_f_name)}')

        return ste.strategy.arg_id, infos

    def run(self, *confs: ConfigType):
        """
        执行策略不同配置
        同时run多个配置时，run必须在__main__下执行
        :param confs: 配置类
        :return:
        """
        self.async_type = len(confs) > 1
        if len(confs) == 0:
            raise PreparationError(f'未输入配置参数')
        elif len(confs) == 1:
            self._stepper(self, confs[0])
        else:
            with ProcessPoolExecutor(max_workers=min(len(confs), os.cpu_count())) as executor:
                result: Dict[str, Dict[str, Any]] = {}
                resp = [executor.submit(self._stepper, self, conf) for conf in confs]
                for future in as_completed(resp):
                    try:
                        res: Tuple[str, Dict[str, Any]] = future.result()
                        result[res[0]] = res[1]
                    except Exception as e:
                        raise PreparationError(f'执行错误：{e.__str__()}')
            df = pd.DataFrame(result).T
            df.to_csv(os.path.join(ROOT_DIR, 'locals', f'batch-{self.strategy_name}.csv'))
